#include "Wmi.h"

#ifdef _WIN32
#ifdef _DEBUG
   #include <crtdbg.h>
   #undef THIS_FILE
   static char THIS_FILE[] = __FILE__;
   #define new       new( _NORMAL_BLOCK, __FILE__, __LINE__)
   #define malloc(s) _malloc_dbg(s, _NORMAL_BLOCK, __FILE__, __LINE__)
#endif
#endif

Wmi::Wmi(void) {
	// Connecting with wbem
	// http://msdn.microsoft.com/en-us/library/windows/desktop/aa389762(v=vs.85).aspx
	// Accessing and manipulating data
	// http://msdn.microsoft.com/en-us/library/windows/desktop/aa389276(v=vs.85).aspx
	// Get WMI data from a remote computer
	// http://msdn.microsoft.com/en-us/library/windows/desktop/aa390422(v=vs.85).aspx
	// ExecQuery method
	// http://msdn.microsoft.com/en-us/library/windows/desktop/aa392107(v=vs.85).aspx

	pSvc = 0;
	pLoc = 0;
}

/**
 * Cleanup
 */
Wmi::~Wmi(void) {
	// pSvc was declared as IWbemServices *pSvc;
	if(pSvc) {
		pSvc->Release();
	}
	// pLoc was declared as IWbemLocator *pLoc;
	if(pLoc) {
		pLoc->Release();     
	}
	//
    CoUninitialize();
}

/**
 * Initialize COM, connecting to localhost
 */
bool Wmi::initialize(void) {
	/**
	 * Initialize COM
	 */
	hr = CoInitializeEx(0,
		COINIT_MULTITHREADED); 
//		COINIT_APARTMENTTHREADED);  // (needed with how we are threading this with each thread having its own Wmi, we should share though.)
	if(FAILED(hr)) {
		cout << "Failed to initialize COM library. Error code = 0x" << hex << hr << endl; 
		return false;
	}

	/**
	 * Set security levels
	 */
	hr = CoInitializeSecurity(
		NULL,							// Security descriptor    
		-1,								// COM negotiates authentication service
		NULL,							// Authentication services
		NULL,							// Reserved
		RPC_C_AUTHN_LEVEL_DEFAULT,		// Default authentication level for proxies
		RPC_C_IMP_LEVEL_IMPERSONATE,	// Default Impersonation level for proxies
		NULL,							// Authentication info
		EOAC_NONE,						// Additional capabilities of the client or server
		NULL);							// Reserved

	if(FAILED(hr) && (hr != RPC_E_TOO_LATE)) {
	   cout << "Failed to initialize security. Error code = 0x" << hex << hr << endl;
	   return false;                  // Program has failed.
	}

	/**
	 * Create a connection to a WMI namespace
	 */
	hr = CoCreateInstance(
		CLSID_WbemLocator,
		0, 
        CLSCTX_INPROC_SERVER,
		IID_IWbemLocator, (LPVOID *) &pLoc);
 
    if(FAILED(hr)) {
        cout << "Failed to create IWbemLocator object. Err code = 0x" << hex << hr << endl;
        return false;     // Program has failed.
    }

	return true;
}

/**
 * Connect to WMI
 */
bool Wmi::connect(void) {
    // Connect to the root\default namespace with the current user.
    hr = pLoc->ConnectServer(
            BSTR(L"root\\cimv2"),
            NULL, NULL, NULL, NULL, 0, NULL, &pSvc);

    if(FAILED(hr)) {
        cout << "Could not connect. Error code = " << hr << endl;
		switch(hr) {
		case WBEM_E_INVALID_PARAMETER:
			break;
		case WBEM_E_OUT_OF_MEMORY:
			break;
		case WBEM_E_UNEXPECTED:
			break;
		case WBEM_E_TRANSPORT_FAILURE:
			break;
		case WBEM_S_NO_ERROR:
			break;
		case WBEM_S_FALSE:
			break;
		case WBEM_S_TIMEDOUT:
			break;
		}

        return false;      // Program has failed.
    }

	return true;
}

/**
 * Initialize COM, connecting to remote server
 */
bool Wmi::connect(char *host, char *user, char *pass) {
	wchar_t pszName[DEFAULT_BUFLEN] = L"losthed\\administrator";
	wchar_t pszPwd[DEFAULT_BUFLEN] = L"putPASShere";

	//wchar_t pszAuthority[DEFAULT_BUFLEN] = L"kERBEROS:losthed";
	//sprintf(pszAuthority, "kERBEROS:%s", "losthed");
	hr = pLoc->ConnectServer(
        _bstr_t(L"\\\\losthed\\root\\cimv2"),
        _bstr_t(pszName),
		_bstr_t(pszPwd),
		NULL,
		NULL,
		NULL,
		NULL,
		&pSvc);
	printf("using credentials\n");

    if(FAILED(hr)) {
        cout << "Could not connect. Error code = " << hr << endl;
		switch(hr) {
		case WBEM_E_INVALID_PARAMETER:
			break;
		case WBEM_E_OUT_OF_MEMORY:
			break;
		case WBEM_E_UNEXPECTED:
			break;
		case WBEM_E_TRANSPORT_FAILURE:
			break;
		case WBEM_S_NO_ERROR:
			break;
		case WBEM_S_FALSE:
			break;
		case WBEM_S_TIMEDOUT:
			break;
		}

        return false;      // Program has failed.
    }

	return true;
}

/**
 * Set security levels
 */
bool Wmi::setSecurity(void) {
    // Set the proxy so that impersonation of the client occurs.
    hres = CoSetProxyBlanket(pSvc,
       RPC_C_AUTHN_WINNT,
       RPC_C_AUTHZ_NONE,
       NULL,
       RPC_C_AUTHN_LEVEL_CALL,
       RPC_C_IMP_LEVEL_IMPERSONATE,
       NULL,
       EOAC_NONE
    );

    if(FAILED(hres)) {
		cout << "Could not set proxy blanket. Error code = 0x" << hex << hres << endl;

		return false;      // Program has failed.
    }

	return true;
}

bool Wmi::ping(char *netbios) {
	bool isPing = false;

//	printf("[>Wmi::ping] '%s' start\n", netbios);

	char query[DEFAULT_BUFLEN];
	sprintf(query, "Select * from Win32_PingStatus WHERE Address=\"%s\"", netbios);
	IEnumWbemClassObject *pEnum;
	hr = pSvc->ExecQuery(L"WQL"
		, _bstr_t(query)
		, WBEM_FLAG_RETURN_IMMEDIATELY | WBEM_FLAG_FORWARD_ONLY
		, NULL
		, &pEnum);

	if(hr != WBEM_S_NO_ERROR) {
		printf("call error: %d\n", hr);
		return false;
	}
//	printf("call succeeded\n");

	HRESULT hRes = WBEM_S_NO_ERROR;
	IWbemClassObject *_next[1];	// or *_next[10]  allocate memory for the found objects
	_next[0] = 0;
	ULONG uReturned;

	VARIANT v;

	// loop through results
//	while(WBEM_S_NO_ERROR == hRes) {
//		printf("inside loop\n");

		// pop off next result
		hRes = pEnum->Next(WBEM_INFINITE, 1, _next, &uReturned);
		// if we got something
		if(SUCCEEDED(hRes)) {
			// do something with the object (this instance is now a local-only copy)
			hr = _next[0]->Get(_bstr_t(L"StatusCode"), 0, &v, 0, 0);
			// check if the call worked
			if(hr == WBEM_S_NO_ERROR) {
				// check the return value to see if it is zero
				//wsprintf(result, L"The status code is %d.\n", V_BSTR(&v));
				printf("StatusCode: %d\n", V_INT(&v));
				if(V_INT(&v) == 0) {
					printf("ping worked, result = true\n");
					VariantClear(&v);
					isPing = true;
				}
				else {
					printf("ping worked, result = failed\n");
					isPing = false;
				}
			}
			else {
				printf("[Wmi::ping]->Get failed\n");
				isPing = false;
			}
		}
//	}
	// release the memory held by the object
	if(_next[0]) {
		_next[0]->Release();
	}

	return isPing;
}